import streamlit as st
import os
import base64
from request_util import convert_image_style, get_image_description
from PIL import Image
from datetime import datetime

# 确保assets文件夹存在
ASSETS_DIR = "assets"
if not os.path.exists(ASSETS_DIR):
    os.makedirs(ASSETS_DIR)

def resize_image(image_path, max_height=400):
    """调整图片尺寸，保持比例，限制最大高度"""
    img = Image.open(image_path)
    # 计算缩放比例
    ratio = max_height / float(img.size[1])
    new_width = int(float(img.size[0]) * ratio)
    # 调整图片尺寸
    img = img.resize((new_width, max_height), Image.Resampling.LANCZOS)
    return img

st.set_page_config(page_title="图片风格转换", layout="wide")
st.title("图片风格转换应用")

# 创建两列布局
col1, col2 = st.columns(2)

# 左侧列：上传原始图片
with col1:
    st.header("原始图片")
    original_uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
    
    if original_uploaded_file is not None:
        # 保存上传的图片到assets文件夹
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] timestamp: {timestamp}")
        original_input_path = os.path.join(ASSETS_DIR, f"temp_input_{timestamp}.png")
        with open(original_input_path, "wb") as f:
            f.write(original_uploaded_file.getvalue())
        # 显示调整尺寸后的原始图片
        resized_img = resize_image(original_input_path)
        st.image(resized_img)

# 右侧列：显示生成结果
with col2:
    st.header("转换结果")
    prompt = st.text_input("请输入转换提示词", value="包裹毛茸茸外表，贴合，软萌，柔和色彩")
    
    if st.button("开始转换") and original_uploaded_file is not None:
        with st.spinner("正在转换中..."):
            # 将图片转换为base64
            with open(original_input_path, "rb") as original_image_file:
                original_base64_image = base64.b64encode(original_image_file.read()).decode('utf-8')

                # 获取图片描述
                response1 = get_image_description(original_base64_image)
                prompt = response1.choices[0].message.content + ", "+ prompt
            
            # 调用转换函数
            print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] prompt: {prompt}")
            
            output_path = os.path.join(ASSETS_DIR, f"output_{timestamp}.png")
            success = convert_image_style(
                inputs=original_base64_image,
                prompt=prompt,
                output_path=output_path
            )
            
            if success:
                st.success("转换成功！")
                # 使用相同的时间戳命名输出文件
                resized_output = resize_image(output_path)
                st.image(resized_output)
            else:
                st.error("转换失败，请重试")
